#include <stdio.h>
#include <gmp.h>
#include "mpzqform.h"

/*
    Copyright 2009-2010 Andrew V. Sutherland

    This file is part of "smoothrelation".

    smoothrelation is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 2 of the License, or
    (at your option) any later version.

    smoothrelation is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with smoothrelation.  If not, see <http://www.gnu.org/licenses/>.
*/

static inline void mpz_mulm (mpz_t o, mpz_t a, mpz_t b, mpz_t p)
	{ mpz_mul(o,a,b); mpz_mod(o,o,p); }

int mpz_qform_valid_discriminant (mpz_t D)
{
	int i;
	
	i = mpz_fdiv_ui(D,4);
	if ( i != 0 && i != 1 ) return 0;
	if ( mpz_perfect_square_p(D) ) return 0;
	return 1;
}

void mpz_qform_discriminant (mpz_t discriminant, mpz_t a, mpz_t b, mpz_t c)
{
	static mpz_t x;
	static int init;

	if ( ! init ) { mpz_init (x);  init = 1; }
	mpz_mul (discriminant, b, b);
	mpz_mul (x, a, c);
	mpz_mul_ui (x, x, 4);
	mpz_sub (discriminant, discriminant, x);
}

void mpz_qform_c2 (mpz_t c, mpz_t a, mpz_t b, mpz_t D)
{
	// compute c = (D-b^2)/(4a)
	mpz_mul (c, b, b);
	mpz_sub (c, c, D);
	mpz_div_2exp (c, c, 2);
	mpz_divexact (c, c, a);
}

void mpz_qform_invert2 (mpz_t a, mpz_t b, mpz_t D)
{
	static int init;
	static mpz_t c;
	
	if ( ! init ) { mpz_init(c); init = 1; }
	if ( mpz_cmp (a,b) != 0 ) {
		mpz_qform_c2 (c, a, b, D);
		if ( mpz_cmp (a, c) != 0 ) mpz_neg (b, b);
	}
}

void mpz_qform_set_id2 (mpz_t a, mpz_t b, mpz_t D)
{
	mpz_set_ui(a,1);
	if ( mpz_divisible_ui_p(D,4) ) mpz_set_ui(b,0); else mpz_set_ui(b,1);
}

int mpz_qform_primeform2 (mpz_t a, mpz_t b, mpz_t p, mpz_t D)
{
	static int init;
	static mpz_t c;
	int i;
	
	if ( ! init ) { mpz_init(c); init = 1; }
	if ( mpz_cmp_ui(p,2) == 0 ) {
//		if ( mpz_divisible_ui_p(D,16) ) gmp_printf ("Warning, primeform with norm %Zd is not prime to the conductor of D=%Zd\n", p, D);
		i = mpz_fdiv_ui(D,8);
		switch (i) {
		case 0: mpz_set_ui(b,0); break;
		case 1: mpz_set_ui(b,1); break;
		case 4: mpz_set_ui(b,2); break;
		default: return 0;
		}
		mpz_set(a,p);
	} else {
		mpz_mul(c,p,p);
//		if ( mpz_divisible_p(D,c) ) gmp_printf ("Warning, primeform with norm %Zd is not prime to the conductor of D=%Zd\n", p, D);
		if ( ! mpz_sqrt_modprime (b,D,p) ) return 0;
		if ( mpz_divisible_ui_p (D,4) ) {
			if ( ! mpz_divisible_ui_p(b,2) ) mpz_sub(b,p,b);					// make b even when D is 0 mod 4
		} else {
			if ( mpz_divisible_ui_p(b,2) ) mpz_sub(b,p,b);						// make b odd when D is 1 mod 4
		}
	}
	mpz_set (a,p);
	mpz_qform_c2 (c, a, b, D);
	if ( mpz_cmp(a,b)==0 && mpz_divisible_p(c,a) ) return 0;					// form not primitive
	mpz_qform_reduce(a,b,c);
	return 1;
}

void mpz_qform_exp2 (mpz_t a2, mpz_t b2, mpz_t a1, mpz_t b1, unsigned long e, mpz_t D, mpz_t L)
{
	static int init;
	static mpz_t u,v,w,c1;
	unsigned long m,n;
	int i;
	
	if ( ! e ) { mpz_qform_set_id2(a2,b2,D); return; }
	if ( ! init ) { mpz_init(u); mpz_init(v); mpz_init(w); mpz_init(c1); init = 1; }
	for ( m=1, n=e>>1; n ; n>>=1, m<<=1 );
	mpz_qform_c2(c1,a1,b1,D);
	mpz_set(u,a1); mpz_set(v,b1); mpz_set(w,c1);
	for ( m >>= 1 ; m ; m >>=1 ) {
		mpz_qform_square(u,v,w,u,v,w,L);
		if ( e&m ) mpz_qform_compose(u,v,w,u,v,w,a1,b1,c1,L);
	}
	mpz_set(a2,u); mpz_set(b2,v);
}

int mpz_qform_compose2 (mpz_t a3, mpz_t b3, mpz_t a1, mpz_t b1, mpz_t a2, mpz_t b2, mpz_t D, mpz_t L, int reduce)
{
	static mpz_t d1, d2, v, w, x, c1, c2, c3;
	static int init;

	if ( ! init ) { mpz_init (d1);  mpz_init (d2);  mpz_init (v);  mpz_init (w);  mpz_init (x); mpz_init (c1); mpz_init (c2);  mpz_init (c3);  init = 1; }

	if ( mpz_cmp(a1,a2)==0 && mpz_cmp(b1,b2)==0 ) return mpz_qform_square2(a3,b3,a1,b1,D,L);
	mpz_qform_c2(c1,a1,b1,D);
	mpz_qform_c2(c2,a2,b2,D);
	if ( mpz_cmp (c1, c2) < 0 ) {
		return mpz_qform_compose (a3, b3, c3, a2, b2, c2, a1, b1, c1, L);
	} else {
		return mpz_qform_compose (a3, b3, c3, a1, b1, c1, a2, b2, c2, L);
	}
}

int mpz_qform_square2 (mpz_t a3, mpz_t b3, mpz_t a1, mpz_t b1, mpz_t D, mpz_t L)
{
	static mpz_t d1, v, w, x, c1, c3;
	static int init;

	if ( ! init ) { mpz_init (d1);  mpz_init (v);  mpz_init (w);  mpz_init (x); mpz_init (c1);  mpz_init (c3); init = 1; }
	mpz_qform_c2 (c1, a1, b1, D);
	return mpz_qform_square (a3, b3, c3, a1, b1, c1, L);
}

// Algorithm 1.5.1 in [CANT] - standard Tonelli-Shanks n^2 algorithm
int mpz_sqrt_modprime (mpz_t o, mpz_t a, mpz_t p)
{
	static mpz_t q, x, y, b;
	static int init;
	int i, r, m;

	if ( ! init ) { mpz_init (q);  mpz_init (x);  mpz_init (y);  mpz_init (b);  init = 1; }
	if ( ! mpz_sgn (a) ) { mpz_set_ui (o, 0);  return 1; }
	if ( mpz_cmp_ui (a, 1) == 0 ) { mpz_set_ui (o, 1);  return 1; }
	if ( mpz_divisible_p(a,p) ) { mpz_set_ui(o,0); return 1; }
	mpz_sub_ui (q, p, 1);
	mpz_set_ui (x, 2);
	r = mpz_remove (q, q, x);
	mpz_set_ui(x,3);
	do {
		mpz_add_ui(x,x,2);
	} while ( mpz_jacobi (x, p) != -1 );
	mpz_powm (y, x, q, p);
	mpz_sub_ui (q, q, 1);
	mpz_divexact_ui (q, q, 2);
	mpz_powm (x, a, q, p);
	mpz_mulm (b, a, x, p);
	mpz_mulm (b, b, x, p);
	mpz_mulm (x, a, x, p);
	for (;;) {
		if ( mpz_cmp_ui (b, 1) == 0 ) { mpz_set (o, x);  break; }
		mpz_set (q, b);
		for ( m = 1 ; m <= r; m++ ) {
			mpz_mulm (q, q, q, p);
			if ( mpz_cmp_ui (q, 1) == 0 ) break;
		}
		if ( m > r ) { gmp_printf ("Unexpected result: m=%d > r=%d in mpz_sqrt_modprime?!  a = %Zd, p = %Zd\n", m, r, a, p);  return 0; }
		if ( m == r ) return 0;
		mpz_set (q, y);
		for ( i = 0 ; i < r-m-1 ; i++ ) mpz_mulm (q, q, q, p);
		mpz_mulm (y, q, q, p);
		r = m;
		mpz_mulm (x, x, q, p);
		mpz_mulm (b, b, y, p);
	}
///	dbg_printf ("%Zd is a square root of %Zd mod %Zd\n", o, a, p);
	return 1;
}

// Form is assumed to be positive definite, primitive, and the discriminant is assumed to be negative
// Algorithm 5.4.2 in Cohen  "A Course in Computational Algebraic Number Theory" p. 243
// sanity checks have been commented our for speed
int mpz_qform_reduce (mpz_t a, mpz_t b, mpz_t c)
{
	static mpz_t x, q, r;
	static int init;
	int i, ambi;

	if ( ! init ) { mpz_init (x);  mpz_init(q);  mpz_init (r);  init = 1; }
//	if ( mpz_sgn (a) <= 0 || mpz_sgn (c) <= 0 ) { err_printf ("Form must be positive definite! a = %Zd, c = %Zd", a, c);  exit (1); }
	ambi = 0;
	for (;;) {
		if ( mpz_cmp (a, c) > 0 ) {
			mpz_set (x, a);
			mpz_set (a, c);
			mpz_set (c, x);
			mpz_neg (b, b);
		}
		i = mpz_cmpabs (b, a);
		if ( i > 0 || (i == 0 && mpz_sgn(b) == -1) ) {
			mpz_mul_2exp (x, a, 1);
			mpz_fdiv_qr (q, r, b, x);
			if ( mpz_cmp (r, a) > 0 ) { mpz_sub (r, r, x); mpz_add_ui (q, q, 1); }
			mpz_add (x, b, r);
			mpz_mul (x, x, q);
			mpz_fdiv_q_2exp (x, x, 1);
			mpz_sub (c, c, x);
			mpz_set (b, r);
		} else {
			if ( i == 0 || ! mpz_sgn(b) ) ambi = 1;
			break;
		}
	}
	if ( mpz_cmp (a, c) == 0 ) {
		if ( mpz_sgn(b) < 0 ) mpz_neg (b, b);
		ambi = 1;
	}
//if ( mpz_cmp (a, c) > 0 ) { puts ("Error in qform_reduce a>c");  exit (0); }
//if ( mpz_cmpabs (b, a) > 0 ) { puts ("Error in qform_reduce |b|>|a|");  exit (0); }
//if ( mpz_sgn (a) <= 0 || mpz_sgn (c) <= 0 ) { err_printf ("Error in qform_reduce, not posdef! a = %Zd, c = %Zd", a, c);  exit (1); }
	return ambi;
}


int mpz_qform_compose (mpz_t u3, mpz_t v3, mpz_t w3, mpz_t u1, mpz_t v1, mpz_t w1, mpz_t u2, mpz_t v2, mpz_t w2, mpz_t L)
{
	static mpz_t Ax, Bx, By, Cy, Dy, F, G, H, Q1, Q2, Q3, Q4, ax, ay, bx, by, cx, cy, dx, dy, m, q, s, b, c, x, y, t, t1, t2, ell;
	static int init;
	unsigned long z;
	
	if ( ! init ) { mpz_init (G); mpz_init (Ax);  mpz_init (Bx); mpz_init (By); mpz_init (F); mpz_init (G); mpz_init (H); mpz_init (m);
					mpz_init (Q1);  mpz_init (Q2);  mpz_init (Q3);  mpz_init (Q4);
					mpz_init (ax); mpz_init (ay);  mpz_init (bx);  mpz_init (by);  mpz_init (dx); mpz_init (cx);  mpz_init (cy);  mpz_init (dy);
					mpz_init (Cy);  mpz_init (Dy);  mpz_init (ell);
					mpz_init (q);  mpz_init (s);  mpz_init (x);  mpz_init (y); mpz_init (b);  mpz_init (c); mpz_init (t1); mpz_init (t2);  init = 1; }
step1:
	mpz_add (s, v1, v2);
	mpz_div_2exp (s, s, 1);
	mpz_sub (m, v2, s);
step2:
    mpz_gcdext (F, b, c, u2, u1);
    if ( mpz_divisible_p (s, F) ) {
    	mpz_set (G, F);
    	mpz_set (Ax, G);
    	mpz_mul (Bx, m, b);
    	mpz_divexact (By, u1, G);
    	mpz_divexact (Cy, u2, G);
    	mpz_divexact (Dy, s, G);
    	goto step5;
    } else {
step3:
		mpz_gcdext (G, 0, y, F, s);
		mpz_divexact (H, F, G);
		mpz_divexact (By, u1, G);
		mpz_divexact (Cy, u2, G);
		mpz_divexact (Dy, s, G);
	}
step4:
	mpz_mod (t1, w1, H);  mpz_mul (t1, t1, b);   mpz_mod (t2, w2, H);  mpz_mul (t2, t2, c);  mpz_add (t1, t1, t2);  mpz_mul (ell, y, t1);  mpz_mod(ell, ell, H);
	mpz_divexact (t1, m, H);  mpz_mul (Bx, b, t1);  mpz_divexact (t1, By, H); mpz_mul (t1, t1, ell);  mpz_add (Bx, Bx, t1);
step5:
	mpz_mod (bx, Bx, By);
	mpz_set (by, By);
	mpz_set_ui (x, 1);  mpz_set_ui (y, 0);  z = 0;
step5b:
	if ( mpz_sgn (bx) && mpz_cmpabs (by, L) > 0 ) goto step5c;
	if ( z & 0x1 ) { mpz_neg (by, by);  mpz_neg (y, y); }
	mpz_mul (ax, G, x);
	mpz_mul (ay, G, y);
	goto step6;
step5c:
	mpz_fdiv_qr (q, t, by, bx);
	mpz_set (by, bx);
	mpz_set (bx, t);
	mpz_mul (t1, q, x);
	mpz_sub (t, y, t1);
	mpz_set (y, x);
	mpz_set (x, t);
	z++;
	goto step5b;
step6:
	if ( z ) goto step7;
	mpz_mul (Q1, Cy, bx);
	mpz_sub (cx, Q1, m);  mpz_divexact (cx, cx, By);
	mpz_mul (dx, bx, Dy);  mpz_sub (dx, dx, w2);  mpz_divexact (dx, dx, By);
	mpz_mul (u3, by, Cy);
	mpz_mul (w3, bx, cx);  mpz_mul (t1, G, dx);  mpz_sub (w3, w3, t1);
	mpz_mul_2exp (t1, Q1, 1);  mpz_sub (v3, v2, t1);
	goto step8;
step7:
	mpz_mul (cx, Cy, bx);  mpz_mul (t1, m, x);  mpz_sub (cx, cx, t1);  mpz_divexact (cx, cx, By);
	mpz_mul (Q1, by, cx);  mpz_add (Q2, Q1, m);
	mpz_mul (dx, Dy, bx);  mpz_mul (t1, w2, x);  mpz_sub (dx, dx, t1);  mpz_divexact (dx, dx, By);
	mpz_mul (Q3, y, dx);  mpz_add (Q4, Q3, Dy);  mpz_divexact (dy, Q4, x);
	if ( mpz_sgn (bx) ) { mpz_divexact (cy, Q2, bx); } else { mpz_mul (cy, cx, dy);  mpz_sub (cy, cy, w1);  mpz_divexact (cy, cy, dx); }
	mpz_mul (u3, by, cy);  mpz_mul (t1, ay, dy);  mpz_sub (u3, u3, t1);
	mpz_mul (w3, bx, cx);  mpz_mul (t1, ax, dx);  mpz_sub (w3, w3, t1);
	mpz_add (v3, Q3, Q4);  mpz_mul (v3, v3, G);  mpz_sub (v3, v3, Q1);  mpz_sub (v3, v3, Q2);
step8:
	return mpz_qform_reduce (u3, v3, w3);	
}

/*
	The algorithm below is from Jacobson and van Der Poorten "Computational Aspects of NUCOMP".
*/
int mpz_qform_square (mpz_t u3, mpz_t v3, mpz_t w3, mpz_t u, mpz_t v, mpz_t w, mpz_t L)
{
	static mpz_t G, Bx, By, Dy, Q1, ax, bx, by, dx, ay, dy, q, x, y, t, t1, t2;
	static int init;
	unsigned long z;
	
	if ( ! init ) { mpz_init (G); mpz_init (Bx); mpz_init (By); mpz_init (Dy); mpz_init (Q1);
					mpz_init (ax);  mpz_init (bx);  mpz_init (by);  mpz_init (dx); mpz_init (ay);  mpz_init (dy);
					mpz_init (q);  mpz_init (x);  mpz_init (y); mpz_init (t1); mpz_init (t2);  init = 1; }

	mpz_gcdext (G, 0, y, u, v);
	mpz_divexact (By, u, G);
	mpz_divexact (Dy, v, G);
	mpz_mul (Bx, y, w); mpz_mod(Bx,Bx,By);
	mpz_set (bx, Bx);
	mpz_set (by, By);
	mpz_set_ui (x, 1);  mpz_set_ui (y, 0);  z = 0;
step_3b:
	if (  mpz_sgn (bx) && mpz_cmpabs(by, L) > 0 ) goto step_3c;
	if ( z&0x1 ) { mpz_neg (by, by);  mpz_neg (y, y);  }
	mpz_mul (ax, G, x);
	mpz_mul (ay, G, y);
	goto step_4;
step_3c:
	mpz_fdiv_qr (q, t, by, bx);
	mpz_set (by, bx);
	mpz_set (bx, t);
	mpz_mul (t1, q, x);
	mpz_sub (t, y, t1);
	mpz_set (y, x);
	mpz_set (x, t);
	z++;
	goto step_3b;
step_4:	
	if ( z ) goto step_5;
	mpz_mul (t1, bx, Dy);
	mpz_sub (t1, t1, w);
	mpz_divexact (dx, t1, By);
	mpz_mul (u3, by, by);
	mpz_mul (w3, bx, bx);
	mpz_add (t1, bx, by);
	mpz_mul (t1, t1, t1);
	mpz_sub (v3, v, t1);
	mpz_add (v3, v3, u3);
	mpz_add (v3, v3, w3);
	mpz_mul (t1, G, dx);
	mpz_sub (w3, w3, t1);
	goto step_6;
step_5:
	mpz_mul (t1, bx, Dy);
	mpz_mul (t2, w, x);
	mpz_sub (t1, t1, t2);
	mpz_divexact (dx, t1, By);
	mpz_mul (Q1, dx, y);
	mpz_add (dy, Q1, Dy);
	mpz_add (t1, dy, Q1);
	mpz_mul (v3, G, t1);
	mpz_divexact (dy, dy, x);
	mpz_mul (u3, by, by);
	mpz_mul (w3, bx, bx);
	mpz_add (t1, bx, by);
	mpz_mul (t1, t1, t1);
	mpz_sub (v3, v3, t1);
	mpz_add (v3, v3, u3);
	mpz_add (v3, v3, w3);
	mpz_mul (t1, ay, dy);
	mpz_sub (u3, u3, t1);
	mpz_mul (t1, ax, dx);
	mpz_sub (w3, w3, t1);
step_6:
	return mpz_qform_reduce (u3, v3, w3);
}

